package org.zxb.calllocalservice.loadbalancer;

import com.alibaba.cloud.commons.lang.StringUtils;
import com.alibaba.cloud.nacos.NacosDiscoveryProperties;
import com.alibaba.cloud.nacos.loadbalancer.NacosLoadBalancer;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.SneakyThrows;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.context.config.annotation.RefreshScope;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.SelectedInstanceCallback;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.http.HttpHeaders;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import reactor.core.publisher.Mono;

import javax.servlet.http.HttpServletRequest;
import java.util.List;
import java.util.Objects;
import java.util.Optional;

/**
 * nacos 负载均衡转发规则；参考 {@link org.springframework.cloud.loadbalancer.core.RoundRobinLoadBalancer}
 * 1. 优先根据 nacos 在 group 中配置的 ip 转发到相同 ip 的服务实例（同源优先调用）
 * 2. 如前端在请求头配置了需要转发的 ip，则根据请求头中的 ip 转发到相同 ip 的服务实例（联调使用）
 * 3. 最后根据 nacos 默认的负载均衡策略转发到服务实例（同集群优先）
 * <p>
 * 注意：需要在转发时将源 ip 添加到请求头（如 nginx 添加 X-Real-IP、RestTemplate 拦截器添加请求头等）：
 * 注意：需要在转发时将源 ip 添加到请求头（如 nginx 添加）：
 *             http
 *      nginx --------> gateway
 *                         |       http(HttpClient)
 *                         ------------------------->  proxy
 *                                                     |       http(RestTemplate)
 *                                                     ----------------------------------> business server
 *                                                                                               |           http(Feign)
 *                                                                                               ---------------------------> other business server
 * @author zxb
 */
@RefreshScope
public class DevLoadBalanceRuleLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    private static final Logger log = LoggerFactory.getLogger(DevLoadBalanceRuleLoadBalancer.class);

    /**
     * 当前服务的名称
     */
    private final String serviceId;

    /**
     * 服务列表
     */
    private ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;

    /**
     * nacos 提供的负载均衡策略
     */
    private final ReactorServiceInstanceLoadBalancer delegate;

    /**
     * 用于前端联调时指定转发到某个实例的 ip
     */
    private final String customIpHeader = "ip";

    public DevLoadBalanceRuleLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                          String serviceId,
                                          NacosDiscoveryProperties nacosDiscoveryProperties) {
        this(serviceInstanceListSupplierProvider, serviceId, nacosDiscoveryProperties,
                new NacosLoadBalancer(serviceInstanceListSupplierProvider, serviceId, nacosDiscoveryProperties));
    }

    public DevLoadBalanceRuleLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                          String serviceId,
                                          NacosDiscoveryProperties nacosDiscoveryProperties,
                                          NacosLoadBalancer delagate) {
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        this.serviceId = serviceId;
        this.delegate = delagate;
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider
                .getIfAvailable(NoopServiceInstanceListSupplier::new);

        return supplier.get(request).next()
                .map(serviceInstances -> processInstanceResponse(request, supplier, serviceInstances));
    }

    private Response<ServiceInstance> processInstanceResponse(Request<?> request, ServiceInstanceListSupplier supplier,
                                                              List<ServiceInstance> serviceInstances) {
        Response<ServiceInstance> serviceInstanceResponse = getInstanceResponse(request, serviceInstances);
        if (supplier instanceof SelectedInstanceCallback && serviceInstanceResponse.hasServer()) {
            ((SelectedInstanceCallback) supplier).selectedServiceInstance(serviceInstanceResponse.getServer());
        }
        return serviceInstanceResponse;
    }

    @SneakyThrows
    private Response<ServiceInstance> getInstanceResponse(Request<?> request, List<ServiceInstance> instances) {
        Response<ServiceInstance> response = new EmptyResponse();
        if (instances.isEmpty()) {
            log.warn("No servers available for service: " + this.serviceId);
            return response;
        }
        // 根据 ip 获取 ServiceInstance
        ServiceInstance serviceInstance = Optional.ofNullable(getServiceInstanceFromRequest(request, instances))
                .orElseGet(() -> getServiceInstanceFromMvc(instances));
        response = new DefaultResponse(serviceInstance);

        // 如果 ip 获取失败，则从委托中获取实例
        if (!response.hasServer() && Objects.nonNull(delegate)) {
            Response[] responseTmp = new Response[]{response};
            delegate.choose(request)
                    .subscribe(resp -> responseTmp[0] = resp);
            response = responseTmp[0];
        }
        return response;
    }

    /**
     * 从 Request 调用中获取请求头的信息，从而获取 ServiceInstance
     *
     * @param request   请求
     * @param instances 服务实例集合
     * @return 符合条件的实例
     */
    public ServiceInstance getServiceInstanceFromRequest(Request<?> request, List<ServiceInstance> instances) {
        Object context = request.getContext();
        // RequestData 类型才能获取到 Headers
        DefaultRequestContext requestContext = null;
        boolean canGetHeaders = context instanceof DefaultRequestContext
                && (requestContext = (DefaultRequestContext) context).getClientRequest() instanceof RequestData;
        if (!canGetHeaders) {
            return null;
        }
        RequestData clientRequest = (RequestData) requestContext.getClientRequest();
        HttpHeaders headers = clientRequest.getHeaders();
        String ipAddress = getIpAddress(headers);

        // 请求的真实 ip 不存在
        if (StringUtils.isBlank(ipAddress)) {
            return null;
        }
        return getServiceInstance(ipAddress, instances);
    }

    /**
     * 从 SpringMvc 的 RequestContextHolder 中获取 ip，从而得到 ServiceInstance
     *
     * @param instances 服务实例集合
     * @return 符合条件的实例
     */
    public ServiceInstance getServiceInstanceFromMvc(List<ServiceInstance> instances) {
        ServletRequestAttributes requestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        // 当异步调用将上下文设置到子线程，将导致 NPE
        if (Objects.isNull(requestAttributes)) {
            return null;
        }
        String ipAddress = getIpAddress(requestAttributes.getRequest());

        // 请求的真实 ip 不存在
        if (StringUtils.isBlank(ipAddress)) {
            return null;
        }
        return getServiceInstance(ipAddress, instances);
    }

    /**
     * 根据 ip 地址获取 ServiceInstance
     *
     * @param ipAddress ip 地址
     * @param instances 服务实例集合
     * @return 符合条件的实例
     */
    private ServiceInstance getServiceInstance(String ipAddress, List<ServiceInstance> instances) {
        if (StringUtils.isBlank(ipAddress)) {
            return null;
        }
        if (log.isTraceEnabled()) {
            log.trace("Transfer to ip: {}", ipAddress);
        }
        return instances.stream()
                .filter(instance -> Objects.equals(ipAddress, instance.getHost())
                        || Objects.equals(ipAddress, instance.getMetadata().get("group"))
                        || Objects.equals(ipAddress, instance.getMetadata().get("nacos.group")))
                // 确保同一 IP 随机调用实例
                .findAny()
                .orElse(null);
    }

    /**
     * 从请求头中获取 ip 地址
     *
     * @param httpServletRequest 请求头
     * @return 从请求头中获取到的 ip 地址
     */
    private String getIpAddress(HttpServletRequest httpServletRequest) {
        // 从自定义的请求头获取 ip，用于前端联调使用
        String ip = httpServletRequest.getHeader(customIpHeader);
        if (StringUtils.isNotBlank(ip)) {
            return ip;
        }
        for (IpHeaderNames headerNameEnum : IpHeaderNames.values()) {
            ip = httpServletRequest.getHeader(headerNameEnum.getName());
            if (StringUtils.isNotBlank(ip) && !"unknown".equalsIgnoreCase(ip)) {
                ip = ip.contains(",") ? ip.split(",")[0] : ip;
                break;
            }
        }
        return ip;
    }

    /**
     * 从请求头中获取 ip 地址
     *
     * @param headers 请求头
     * @return 从请求头中获取到的 ip 地址
     */
    private String getIpAddress(HttpHeaders headers) {
        // 从自定义的请求头获取 ip，用于前端联调使用
        String ip = headers.getFirst(customIpHeader);
        if (StringUtils.isNotBlank(ip)) {
            return ip;
        }
        for (IpHeaderNames headerNameEnum : IpHeaderNames.values()) {
            ip = headers.getFirst(headerNameEnum.getName());
            if (StringUtils.isNotBlank(ip) && !"unknown".equalsIgnoreCase(ip)) {
                ip = ip.contains(",") ? ip.split(",")[0] : ip;
                break;
            }
        }
        return ip;
    }


    @AllArgsConstructor
    @Getter
    private enum IpHeaderNames {
        /**
         * 能够获取源 ip 的请求头
         */

        X_FORWARDED_FOR("x-forwarded-for"),
        X_REAL_IP("X-Real-IP"),
        HTTP_CLIENT_IP("HTTP_CLIENT_IP"),
        PROXY_CLIENT_IP("Proxy-Client-IP"),
        WL_PROXY_CLIENT_IP("WL-Proxy-Client-IP"),
        HTTP_X_FORWARDED_FOR("HTTP_X_FORWARDED_FOR");

        private final String name;
    }
}
